#include <iostream>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <utility>
#include <map>

#include "TH1D.h"
#include "TFile.h"
#include "TTree.h"
#include "TChain.h"
#include "TGraph.h"
#include "TParameter.h"
#include "DelphesClasses.h"
#include "ExRootTreeReader.h"
#include "mr.h"
#include "TClonesArray.h"
#include "covariance.h"
#include "vector_tools.h"
#include "TProfile.h"

#include "CmdLine/CmdLine.hh"
using namespace std;

bool debug=false;

const double PI=3.1415926;

class TLorentzSorter{

public:
  bool operator()(const TLorentzVector& a, const TLorentzVector& b) const
  {
    return a.Pt()>b.Pt();
  }
  


};

double fixphi(double input)
{
  if(input < 0)
    input = -input;

  while(input > 2*PI)
    {
      input -= 2*PI;
    }
  return input;
}

//send deltaphi = PI to infinity
double transformphi(double input)
{
  return log(tan(fixphi(input)/2));
}


//declare a tree class
class MyTree{
public:
   TChain         *t;   //!pointer to the analyzed TTree or TChain

   // Declaration of leaf types
  double weight_scale;

   Float_t         weight;
   Float_t         MET;
   Float_t         MET_phi;
   Float_t         HTjets;
   Float_t         HTleptons;
   Float_t         HTleptons_lep20;
   Float_t         m_eff;
   Float_t         METsignif;
   Int_t           nBtags;
   Int_t           nJets30;
   Int_t           nJets60;
   Int_t           nJets100;
   Int_t           nJets200;
   Float_t         pTJet[100];
   Float_t         etaJet[100];
   Float_t         phiJet[100];
   Float_t         massJet[100];
   UInt_t          btagJet[100];

   Float_t         Tau1Jet[100];
   Float_t         Tau2Jet[100];
   Float_t         Tau3Jet[100];
   Int_t           WTagJet[100];
   Float_t         TrimmedMassJet[100];
   Int_t           TopTagJet[100];

   Float_t         mindPhiMetJet30;
   Float_t         mindPhiMetJet60;
   Float_t         mindPhiMetJet100;
   Float_t         mindPhiMetJet200;
   Float_t         dPhiMETBtag0;
   Float_t         dPhiMETBtag1;
   Float_t         dPhiMETBtag2;
   Float_t         mindPhiMETBtag;
   Int_t           nLeps35;
   Float_t         pTLep[10];
   Float_t         etaLep[10];
   Float_t         phiLep[10];
   Int_t           flavLep[10];
   Float_t         pTTruth[50];
   Float_t         etaTruth[100];
   Float_t         phiTruth[100];
   Float_t         massTruth[100];
   Int_t           pdgIdTruth[100];
   Int_t           M1Truth[100];
   Int_t           M2Truth[100];
   Int_t           D1Truth[100];
   Int_t           D2Truth[100];
   Int_t           statusTruth[100];
  


  MyTree(string input)
  {
    // Set branch addresses and branch pointers
    t = new TChain("OutputTree");

    t->AddFile(input.c_str());
    assert(t->GetFile() != NULL);

    t->SetMakeClass(1);
    
    t->SetBranchAddress("weight", &weight);
    t->SetBranchAddress("MET", &MET);
    t->SetBranchAddress("MET_phi", &MET_phi);
    t->SetBranchAddress("HTjets", &HTjets);
    t->SetBranchAddress("HTleptons", &HTleptons);
    t->SetBranchAddress("HTleptons_lep20", &HTleptons_lep20);
    t->SetBranchAddress("m_eff", &m_eff);
    t->SetBranchAddress("METsignif", &METsignif);
    t->SetBranchAddress("nBtags", &nBtags);
    t->SetBranchAddress("nJets30", &nJets30);
    t->SetBranchAddress("nJets60", &nJets60);
    t->SetBranchAddress("nJets100", &nJets100);
    t->SetBranchAddress("nJets200", &nJets200);
    t->SetBranchAddress("pTJet", pTJet);
    t->SetBranchAddress("etaJet", etaJet);
    t->SetBranchAddress("phiJet", phiJet);
    t->SetBranchAddress("massJet", massJet);
    t->SetBranchAddress("btagJet", btagJet);
    t->SetBranchAddress("Tau1Jet", Tau1Jet);
    t->SetBranchAddress("Tau2Jet", Tau2Jet);
    t->SetBranchAddress("Tau3Jet", Tau3Jet);
    t->SetBranchAddress("WTagJet", WTagJet);
    t->SetBranchAddress("TrimmedMassJet", TrimmedMassJet);
    t->SetBranchAddress("TopTagJet", TopTagJet);

    t->SetBranchAddress("mindPhiMetJet30", &mindPhiMetJet30);
    t->SetBranchAddress("mindPhiMetJet60", &mindPhiMetJet60);
    t->SetBranchAddress("mindPhiMetJet100", &mindPhiMetJet100);
    t->SetBranchAddress("mindPhiMetJet200", &mindPhiMetJet200);
    t->SetBranchAddress("dPhiMETBtag0", &dPhiMETBtag0);
    t->SetBranchAddress("dPhiMETBtag1", &dPhiMETBtag1);
    t->SetBranchAddress("dPhiMETBtag2", &dPhiMETBtag2);
    t->SetBranchAddress("mindPhiMETBtag", &mindPhiMETBtag);
    t->SetBranchAddress("nLeps35", &nLeps35);
    t->SetBranchAddress("pTLep", pTLep);
    t->SetBranchAddress("etaLep", etaLep);
    t->SetBranchAddress("phiLep", phiLep);
    t->SetBranchAddress("flavLep", flavLep);
    t->SetBranchAddress("pTTruth", pTTruth);
    t->SetBranchAddress("etaTruth", etaTruth);
    t->SetBranchAddress("phiTruth", phiTruth);
    t->SetBranchAddress("massTruth", massTruth);
    t->SetBranchAddress("pdgIdTruth", pdgIdTruth);
    t->SetBranchAddress("M1Truth", M1Truth);
    t->SetBranchAddress("M2Truth", M2Truth);
    t->SetBranchAddress("D1Truth", D1Truth);
    t->SetBranchAddress("D2Truth", D2Truth);
    t->SetBranchAddress("statusTruth", statusTruth);


    //now if we have signal, need to change
    weight_scale=-1.0;
    //this is signal file
    if(input.find("decoupledGO")!=string::npos)
      {
	//get the cxn
	TParameter<double>* par=
	  (TParameter<double>*) t->GetFile()->Get("TotalProcessEvents");
	//divide by the total number of generated events
	double nevt=par->GetVal();


	//get a TGraph for the cxn
	double mass[]={50,200,400,600,800,1000,1200,1400,1600,1800,2000,2200,2400,2600,3400,3600,3800,4000,4200,4400,4600,4800,5000,5200,5400,5600,5800,6000,6200,6400,6600,6800,7000,7200,7400,7600,7800,8000,8200,8400,8600,8800,9000,9200,9400,9600,9800,10000};

	double cxn[]={4.180000e+05, 3.000000e+03, 1.950000e+02, 3.660000e+01, 1.020000e+01, 3.960000e+00, 1.810000e+00, 9.450000e-01, 5.210000e-01, 2.910000e-01, 1.794000e-01, 1.090000e-01, 6.860000e-02, 4.440000e-02, 9.880000e-03, 7.100000e-03, 5.160000e-03, 3.800000e-03, 2.840000e-03, 2.140000e-03, 1.622000e-03, 1.244000e-03, 9.600000e-04, 7.460000e-04, 5.840000e-04, 4.600000e-04, 3.660000e-04, 2.900000e-04, 2.320000e-04, 1.872000e-04, 1.512000e-04, 1.226000e-04, 9.980000e-05, 8.140000e-05, 6.680000e-05, 5.480000e-05, 4.520000e-05, 3.740000e-05, 3.100000e-05, 2.580000e-05, 2.140000e-05, 1.788000e-05, 1.496000e-05, 1.254000e-05, 1.052000e-05, 8.840000e-06, 7.460000e-06, 6.280000e-06};

	TGraph g(sizeof(mass)/sizeof(double), mass, cxn);
	
	//now get the mass
	int pos=input.find("Up.T1");
	
	string sub_mass=input.substr(pos+5, 5);
	cout<<"found mass string for stop production: "<<sub_mass<<endl;

	int real_mass = atoi(sub_mass.c_str());

	double real_cxn = g.Eval(real_mass);
	cout<<"cxn: "<<real_cxn<<endl;

	weight_scale= real_cxn * 1e3 / nevt;

	
	cout<<"squark mass: "<<real_mass<<endl;
	cout<<"weight: "<<weight_scale<<endl;
      }

  }


  int cout_truth() const
  {
    for(int i=0; i<50; i++)
      cout<<i<<": \t"<<pdgIdTruth[i]
	  <<"\t"<<M1Truth[i]-2
	  <<"\t"<<statusTruth[i]<<endl;
  }

  //function to find parent of muon
  int mu_parent(int i) const
  {
    if(i < 0 || i >= 50)
      return -1;

    int parent=M1Truth[i]-2;
    
    while(parent >= 0 && parent <50 && 
	  abs(pdgIdTruth[parent])==13 )
      {
	int temp_parent=parent;
	parent=M1Truth[parent]-2; 
	
	//infinite loop
	if(temp_parent == parent)
	  break;
      }

    if(parent >= 0 && parent <50)
      return abs(pdgIdTruth[parent]);

    else
      return -1;
  }

};



int main(int argc, char** argv)
{
  //initialize input
  CmdLine cmdline(argc, argv);

  //get input and output
  string input, output;

  double cut_MR, cut_met;
  
  int cout_n=-1;

  try
    {
      input = cmdline.value<string> ("-i");
      output = cmdline.value<string> ("-o");
      cout_n = cmdline.value<int> ("-n", -1);
    }
  catch(...)
    {
      cout<<"ERROR: please specify input and output via -i input -o output"<<endl;
      return 0;
    }

  //initialize plots
  map<string, TH1D*> hist;
  map<string, TProfile*> profile;

  //book histograms

  hist["cutflow"]=new TH1D("cutflow", "Cut Flow", 10, 0, 10);

  hist["nmu"]=new TH1D("nmu", "number of #mu", 10, 0, 10);

  hist["nmu_100"]=new TH1D("nmu_100", "number of #mu > 100 GeV", 10, 0, 10);

  hist["nmu_500"]=new TH1D("nmu_500", "number of #mu > 500 GeV", 10, 0, 10);

  hist["nmu_1000"]=new TH1D("nmu_1000", "number of #mu > 1000 GeV", 10, 0, 10);

  hist["nmu_100_j100"]=new TH1D("nmu_100_j100", "number of #mu > 100 GeV within 100 GeV jet", 10, 0, 10);

  hist["nmu_500_j100"]=new TH1D("nmu_500_j100", "number of #mu > 500 GeV within 100 GeV jet", 10, 0, 10);

  hist["nmu_1000_j100"]=new TH1D("nmu_1000_j100", "number of #mu > 1000 GeV within 100 GeV jet", 10, 0, 10);

  hist["nmu_100_j500"]=new TH1D("nmu_100_j500", "number of #mu > 100 GeV within 500 GeV jet", 10, 0, 10);

  hist["nmu_500_j500"]=new TH1D("nmu_500_j500", "number of #mu > 500 GeV within 500 GeV jet", 10, 0, 10);

  hist["nmu_1000_j500"]=new TH1D("nmu_1000_j500", "number of #mu > 1000 GeV within 500 GeV jet", 10, 0, 10);

  hist["nmu_100_j1000"]=new TH1D("nmu_100_j1000", "number of #mu > 100 GeV within 1 TeV jet", 10, 0, 10);

  hist["nmu_500_j1000"]=new TH1D("nmu_500_j1000", "number of #mu > 500 GeV within 1 TeV jet", 10, 0, 10);

  hist["nmu_1000_j1000"]=new TH1D("nmu_1000_j1000", "number of #mu > 1000 GeV within 1 TeV jet", 10, 0, 10);

  hist["nmu_100_3"]=new TH1D("nmu_100_3", "number of status 3 #mu > 100 GeV", 10, 0, 10);

  hist["nmu_500_3"]=new TH1D("nmu_500_3", "number of status 3 #mu > 100 GeV", 10, 0, 10);

  hist["nmu_1000_3"]=new TH1D("nmu_1000_3", "number of status 3 #mu > 1000 GeV", 10, 0, 10);

  hist["nmu_100_status"]=new TH1D("nmu_100_status", "status for #mu > 100 GeV", 5, 0, 5);

  hist["nmu_500_status"]=new TH1D("nmu_500_status", "status for #mu > 100 GeV", 5, 0, 5);

  hist["nmu_1000_status"]=new TH1D("nmu_1000_status", "status for #mu > 1000 GeV", 5, 0, 5);


  profile["status3_pt"]=new TProfile("status3_pt", 
				     "Proportion of status 3 muons vs. #it{p_{T}}", 60, 0, 1500);

  profile["status2_pt"]=new TProfile("status2_pt", 
				     "Proportion of status 2 muons vs. #it{p_{T}}", 60, 0, 1500);

  profile["status1_pt"]=new TProfile("status1_pt", 
				     "Proportion of status 1 muons vs. #it{p_{T}}", 60, 0, 1500);

  
  double weight_scale=1.0;
  
  float min_pt=100.;

  //initialize tree reading
  MyTree c(input);

  cout<<"Processing "<<c.t->GetEntries()<<" events..."<<endl;
  int division=ceil(c.t->GetEntries()/20.);
  
  if(cout_n>=0)
    {
      assert(cout_n< c.t->GetEntries());
      c.t->GetEntry(cout_n);
      c.cout_truth();
      return 1;
    }
  
 
  for(Long64_t i=0; i< c.t->GetEntries(); i++)
    {

      if(i% division == 0)
	cout<<(i/division)*5<<" % completed"<<endl;

      if(debug)
	cout<<"at event: "<<i<<endl;

      c.t->GetEntry(i);

      int nj100=c.nJets100;

      if(nj100 < 2 || 
	 fabs(c.etaJet[0]) > 2.5 ||
	 fabs(c.etaJet[1]) > 2.5)
	continue;

      //get jet vectors
      TLorentzVector j1, j2;
      j1.SetPtEtaPhiM(c.pTJet[0], c.etaJet[0], 
		      c.phiJet[0], c.massJet[0]);

      j2.SetPtEtaPhiM(c.pTJet[1], c.etaJet[1], 
		      c.phiJet[1], c.massJet[1]);


      double weight=c.weight;

      //if scale greater than zero
      //it means this is signal, and weight 
      //is computed from weight_scale
      if(c.weight_scale > 0)
	weight = c.weight_scale;

      //quick nmuons and Ws
      int nmu=0;

      int nmu_100 =0 ;
      int nmu_500 =0 ;
      int nmu_1000 =0 ;

      int nmu_100_3 =0 ;
      int nmu_500_3 =0 ;
      int nmu_1000_3 =0 ;

      
      vector<TLorentzVector> mu_list;


      //loop over the list twice
      for(int ic=0; ic<50*2; ++ic)
	{
	  int i=ic;
	  if(ic>=50) 
	    i-=50;

	  if(abs(c.pdgIdTruth[i])!=13) continue;
	  if(c.pTTruth[i]<=0) continue;


	  //fill muon info
	  //test pt
	  if(ic<50)
	    {
	      if(c.pTTruth[i] > 100.)
		{
		  nmu_100 ++;
		  
		  hist["nmu_100_status"]->
		    Fill(c.statusTruth[i], weight);
		  
		  if(c.statusTruth[i] == 3)
		    nmu_100_3++;
		}

	      if(c.pTTruth[i] > 500.)
		{
		  nmu_500 ++;
		  
		  hist["nmu_500_status"]->
		    Fill(c.statusTruth[i], weight);
		  
		  if(c.statusTruth[i] == 3)
		    nmu_500_3++;
		}
	      
	      if(c.pTTruth[i] > 1000.)
		{
		  nmu_1000 ++;
		  
		  hist["nmu_1000_status"]->
		    Fill(c.statusTruth[i], weight);
		  
		  if(c.statusTruth[i] == 3)
		    nmu_1000_3++;
		}
	    }


	  if(fabs(c.etaTruth[i])>2.5) continue;

	  //first loop is over muons not status 2
	  if(ic<50 && c.statusTruth[ic]==2)
	    continue;

	  //second loop is over muons status 2
	  if(ic>=50 && c.statusTruth[i]!=2)
	    continue;
	  
	  TLorentzVector muj;

	  muj.SetPtEtaPhiM(c.pTTruth[i], c.etaTruth[i], 
			   c.phiTruth[i], 0);

	  //check for duplicates
	  bool b_dup=false;
	  for(int k=0; k<mu_list.size(); ++k)
	    {
	      if(mu_list[k].DeltaR(muj) < 1e-3)
		{
		  b_dup=true;
		  break;
		}
	    }
	  if(b_dup)
	    continue;

	  ++nmu;
	  mu_list.push_back(muj);

	  if(c.statusTruth[i]==1)
	    {
	      profile["status1_pt"]->Fill(c.pTTruth[i], 1.);
	      profile["status2_pt"]->Fill(c.pTTruth[i], 0.);
	      profile["status3_pt"]->Fill(c.pTTruth[i], 0.);
	    }	  

	  if(c.statusTruth[i]==2)
	    {
	      profile["status1_pt"]->Fill(c.pTTruth[i], 0.);
	      profile["status2_pt"]->Fill(c.pTTruth[i], 1.);
	      profile["status3_pt"]->Fill(c.pTTruth[i], 0.);
	    }	  

	  if(c.statusTruth[i]==3)
	    {
	      profile["status1_pt"]->Fill(c.pTTruth[i], 0.);
	      profile["status2_pt"]->Fill(c.pTTruth[i], 0.);
	      profile["status3_pt"]->Fill(c.pTTruth[i], 1.);
	    }	  
	}


      sort(mu_list.begin(), mu_list.end(), TLorentzSorter());


      int nmu_100_j100=0;
      int nmu_500_j100=0;
      int nmu_1000_j100=0;

      int nmu_100_j500=0;
      int nmu_500_j500=0;
      int nmu_1000_j500=0;

      int nmu_100_j1000=0;
      int nmu_500_j1000=0;
      int nmu_1000_j1000=0;

      //loop over muons done
      //now loop over jets
      for(int i=0; i<mu_list.size(); ++i)
      for(int ijet=0; ijet<10; ijet++)
	if(fabs(c.etaJet[ijet])<2.5)
	  {
	    TLorentzVector j;
	    j.SetPtEtaPhiM(c.pTJet[ijet], c.etaJet[ijet], 
			   c.phiJet[ijet], c.massJet[ijet]);
	    
	    if(j.DeltaR(mu_list[i]) > 0.5)
	      continue;

	    if(c.pTJet[ijet]>=100.)
	      {
		if(mu_list[i].Pt() > 100)
		  ++nmu_100_j100;

		if(mu_list[i].Pt() > 500)
		  ++nmu_500_j100;
		
		if(mu_list[i].Pt() > 1000)
		  ++nmu_1000_j100;
	      }

	    if(c.pTJet[ijet]>=500.)
	      {
		if(mu_list[i].Pt() > 100)
		  ++nmu_100_j500;

		if(mu_list[i].Pt() > 500)
		  ++nmu_500_j500;
		
		if(mu_list[i].Pt() > 1000)
		  ++nmu_1000_j500;
	      }

	    if(c.pTJet[ijet]>=1000.)
	      {
		if(mu_list[i].Pt() > 100)
		  ++nmu_100_j1000;

		if(mu_list[i].Pt() > 500)
		  ++nmu_500_j1000;
		
		if(mu_list[i].Pt() > 1000)
		  ++nmu_1000_j1000;
	      }
	    break;
	  }


      hist["nmu"]->Fill(nmu, weight);

      hist["nmu_100"]->Fill(nmu_100, weight);
      hist["nmu_500"]->Fill(nmu_500, weight);
      hist["nmu_1000"]->Fill(nmu_1000, weight);

      hist["nmu_100_3"]->Fill(nmu_100_3, weight);
      hist["nmu_500_3"]->Fill(nmu_500_3, weight);
      hist["nmu_1000_3"]->Fill(nmu_1000_3, weight);

      hist["nmu_100_j100"]->Fill(nmu_100_j100, weight);
      hist["nmu_500_j100"]->Fill(nmu_500_j100, weight);
      hist["nmu_1000_j100"]->Fill(nmu_1000_j100, weight);

      hist["nmu_100_j500"]->Fill(nmu_100_j500, weight);
      hist["nmu_500_j500"]->Fill(nmu_500_j500, weight);
      hist["nmu_1000_j500"]->Fill(nmu_1000_j500, weight);

      hist["nmu_100_j1000"]->Fill(nmu_100_j1000, weight);
      hist["nmu_500_j1000"]->Fill(nmu_500_j1000, weight);
      hist["nmu_1000_j1000"]->Fill(nmu_1000_j1000, weight);
    }
  cout<<"100 % completed"<<endl;


  TFile fout(output.c_str(), "RECREATE");
  map<string, TH1D*>::iterator it=hist.begin();
  for(;it!=hist.end(); ++it)
    {
      it->second->Write();
    }

  map<string, TProfile*>::iterator ip=profile.begin();
  for(;ip!=profile.end(); ++ip)
    {
      ip->second->Write();
    }

  fout.Close();
  return 1;
}
